import scipy.io as sio
from qutip import *
from pylab import *
from numpy import *
from math import *
from datetime import *
from nv_system import *


        
class MWChannel:
    def __init__(self,freq,init_phase,field,eff_momentums,phi=0,thi=0):
        self.frequency = freq
        self.phase = init_phase
        self.field = field
        self.phi = phi
        self.theta = thi
        Ix,Iy,Iz = eff_momentums
        self.control_ham = -(Ix*cos(self.theta)*cos(self.phi) + \
            Iy*cos(self.theta)*sin(self.phi)+\
            Iz*sin(self.theta))          
    
class pulse_shape:
    def get_shape(self,channel,t):
        pass
            
    def get_duration(self):
        pass
      
class single_channel_pulse(pulse_shape):   
    def __init__(self,channame,duration=0,ampratio = 1.0,phase = 0.0):
        self.__duration = duration
        self.__amp = ampratio
        self.__phase = phase
        self.channel = channame
    
    def get_shape(self,channel,t):
        if channel!=self.channel:
            return 0,0
        if (t < self.__duration and t >= 0) or self.__duration == 0:
            return self.__amp,self.__phase
        else:
            return 0,0
            
    def get_duration(self):
        return self.__duration

class shaped_pulse:
    def __init__(self,waveforms,slot_width):
        #waveforms in format [[channelname1,x1,y1]...]
        self.ampls={}
        self.phas={}
        self.durs={}
        self.slot_width = slot_width
        if not isinstance(waveforms[0],list):
            waveforms = [waveforms]
            
        for waveform in waveforms:
            chname,x,y=waveform
            self.ampls[chname]=[]
            self.phas[chname]=[]
            self.durs[chname]=float(slot_width*len(x))
            print 'duration of channel ',chname,' is ',self.durs[chname],' ns'
            if not isinstance(x,list):
                x = [x]
                y = [y]
            for i in range(len(x)):
                c = x[i]+1j*y[i]
                self.ampls[chname].append(norm(c))
                self.phas[chname].append(angle(c))
    
        self.__duration = max(self.durs.values())
            
    def get_shape(self,channel,t):
#         print t
        if channel in self.ampls:
            if (t < self.durs[channel] and t >= 0):
                n = int(floor(t/self.slot_width))
                return self.ampls[channel][n],self.phas[channel][n]
            else:
                return 0.0,0.0
            
    def get_duration(self):
        return self.__duration

# class gaussian_pulse:
#     def __init__(self,a,b,c):
#         self.a = a
#         self.b = b
#         self.c = c
#         
#     def cal_gaussian_function(self):
#         return self.a*exp(-(x-b)**2/c**2/2)
#     
#     def get_shape(self,channel,t):
#         if not channel in self.ampls:
#             return 0,0
#         if t > self.durs[channel]:
#             return 0,0
#         n = int(floor(t / self.slot_width))
#         return self.ampls[channel][n],self.phas[channel][n]
#                 

class multi_channel_pulse(pulse_shape):
    def __init__(self,channames,duration,ampratio = 1.0,phase = 0.0):
        self.__duration = duration
        self.channels = channames
        if not isinstance(ampratio,list):
            ampratio = [ampratio]*len(self.channels)
        if not isinstance(phase,list):
            phase = [phase]*len(self.channels)
        self.__amp = ampratio
        self.__phase = phase
    
    def get_shape(self,channel,t):
        if channel in self.channels:
            if (t < self.__duration and t >= 0) or self.get_duration() == 0:
                i = self.channels.index(channel)
                return self.__amp[i],self.__phase[i]
            else:
                return 0.0,0.0
            
    def get_duration(self):
        return self.__duration 
       
# class pulse_slice:
#     def __init__(self,*args):
#         self.shapes = {}
#         #args = chname1,pulse1,chname2,pulse2...target_operator
#         for i in range(len(args)/2):
#             self.shapes[args[i*2]]=args[i*2+1]
#         self.duration = 0
#         for pshape in self.shapes.values():
#             self.duration = max(self.duration,pshape.get_duration())
#             
#         if len(args) % 2 == 1:
#             self.target_operator = args[len(args)-1]
#         else:
#             self.target_operator = None
#     
#     def get_shape(self,ch,t):
#         return self.shapes[ch].get_shape(t)
#         
#     def get_duration(self):
#         return self.duration
        
# class pulse_controller:
#     def __init__(self,pulse=None):
#         self.times = [0]
#         self.pointer = 0
#         self.sequence = []
#         if pulse:
#             self.append(pulse)
#         
#     def append(self,pthing):
#         if isinstance(pthing,list) :
#             for pulse in pthing:
#                 self.sequence.append(pulse)               
#         elif isinstance(pthing,pulse_shape):
#             self.sequence.append(pthing)       
#             self.times.append(self.times[-1] + pthing.get_duration())
#             
#     def get_total_duration(self):
#         tduration = 0
#         for pulse in self.sequence:
#             tduration = tduration + pulse.duration
#         return tduration
#     
#     def get_shape(self,t):
#         if len(self.sequence) == 1 and self.sequence[0].get_duration()==0:
#             return self.sequence[0].get_shape(0)
#         
#         while t <= self.times[self.pointer]:
#             self.pointer = self.pointer - 1
#         while t > self.times[self.pointer] + self.sequence[self.pointer].get_duration() and self.pointer < len(self.times):
#             self.pointer = self.pointer + 1
#             
#         if self.pointer == len(self.times):
#             return 0.0,0.0
#         else:
#             return self.sequence[self.pointer].get_shape(t-self.times[self.pointer])
#             
            